#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_update_resist_gene.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/31 11:12
@Description : mNGS 流程上传耐药基因检出结果到数据库
"""

import logging
import re
from argparse import ArgumentParser, Namespace
from pathlib import Path
from subprocess import run, PIPE

# --------
# 设置日志配置
# --------
logging.basicConfig(
    level=logging.INFO,  # 设置日志级别为 INFO
    format="[%(asctime)s] %(levelname)s: %(message)s",  # 日志格式
    datefmt="%Y-%m-%d %H:%M:%S",  # 日期时间格式
)

# 创建日志记录器
LOGGER = logging.getLogger("mNGS")

# --------
# 配置常量
# --------
# 工作目录，包含数据和结果
WORK_DIR = Path("/data/mNGS/runmngs/result/")

# 脚本路径，使用该脚本上传数据
SCRIPT = Path("/data/mNGS/pipeline/script/autoReport/biotools/main.py")

# Python 3 解释器路径
PYTHON3 = Path("/home/bioinfo/software/miniconda3/envs/python-venv/bin/python3")

# CSV 文件存储目录，包含租户 ID 文件
CSV_DIR = Path("/data/mNGS/runmngs/csv/")


# --------
# 辅助函数
# --------
def _get_tenant_id(jobid: str) -> int:
    """
    获取当前任务的租户编号
    :param jobid: 任务的 jobid，用于查找租户 ID 文件
    :return: 租户编号，如果无法获取则返回默认值 1
    """
    id_file = CSV_DIR / f"{jobid}.id"  # 租户 ID 文件路径

    # 检查文件是否存在，且文件大小非零
    if not id_file.exists() or id_file.stat().st_size == 0:
        LOGGER.warning(f"租户文件 {id_file} 不存在或为空，返回默认租户编号 1")
        return 1

    # 尝试读取租户 ID
    try:
        with open(id_file, "r") as f:
            tenant_id = f.read().strip()  # 去除多余的空白字符
            # 如果读取内容不是数字，返回默认值
            if not tenant_id.isdigit():
                LOGGER.warning(f"租户编号 {tenant_id} 非法，返回默认值 1")
                return 1
            return int(tenant_id)
    except (OSError, ValueError) as err:
        LOGGER.error(f"读取租户文件 {id_file} 时出错: {err}，返回默认租户编号 1")
        return 1


def _run_or_die(cmd: str):
    """
    执行命令，如果命令失败则记录错误并抛出异常。
    :param cmd: 要执行的命令字符串
    """
    try:
        LOGGER.info(f"执行命令：{cmd}")
        result = run(
            cmd,
            shell=True,  # 允许通过 shell 执行命令
            text=True,  # 返回文本格式的输出
            stdout=PIPE,  # 捕获标准输出
            stderr=PIPE,  # 捕获标准错误输出
        )

        # 打印命令的标准输出和标准错误
        if result.stdout:
            LOGGER.info(f"标准输出：\n{result.stdout}")
        if result.stderr:
            LOGGER.error(f"标准错误：\n{result.stderr}")

        # 检查命令执行结果
        if result.returncode != 0:
            LOGGER.error(f"命令失败：{cmd}，退出码 {result.returncode}")
            raise RuntimeError(f"命令执行失败，退出码 {result.returncode}")

        LOGGER.info(f"命令完成：{cmd}")

    except Exception as e:
        LOGGER.error(f"命令失败：{cmd}，错误：{e}")
        raise  # 抛出异常，确保错误得到传递


# --------
# 主函数
# --------
def main(args: Namespace) -> None:
    """
    主函数，获取 job_id 和租户 ID，处理输入文件并调用脚本进行数据上传。
    :param args: 命令行参数
    """
    job_id = args.job_id
    tenant_id = _get_tenant_id(job_id)  # 获取租户编号

    # 输入目录，包含每个样本的基因映射数据
    in_dir = WORK_DIR / job_id / job_id / "rgi"

    # 获取目录下所有文件，筛选出基因映射文件
    in_files = list(in_dir.glob("*.gene_mapping_data.en.txt"))

    # 获取样本名称（去掉文件后缀）
    samples = [
        re.sub(r"\.gene_mapping_data\.en\.txt$", "", in_file.name)
        for in_file in in_files
    ]

    # 对每个文件和样本执行更新操作
    for in_file, sample in zip(in_files, samples):
        # 构建上传命令
        command = (
            f"{PYTHON3} {SCRIPT} update_resist_gene "
            f"--in_file {in_file} "
            f"--sample_id {sample} "
            f"--job_id {job_id}  "
            f"--tenant_id {tenant_id}"
        )
        _run_or_die(command)  # 执行命令


# --------
# 脚本入口
# --------
if __name__ == "__main__":
    # 解析命令行参数
    parser = ArgumentParser(description="mNGS 流程上传耐药基因检出结果到数据库")
    parser.add_argument(
        "--job_id",  # 芯片编号，用于唯一标识任务
        type=str,
        required=True,
        help="芯片编号",
    )
    # 执行主函数
    main(args=parser.parse_args())
